#include "winograd_36_to_4x4.h"

namespace mslite {

Winograd36To4x4::Winograd36To4x4(cl_context context, cl_device_id device, std::string clKernelName, std::vector<int> ishape, std::vector<int> oshape): Operation(context) {
    this->clFileName = clKernelName + ".cl";
    this->clKernelName = clKernelName;
    this->context = context;
    this->inputNames = {"input", "bias"};
    std::vector<int> bshape(1, ishape[3]);
    this->inputShapes = {ishape, bshape};                                   // ishape must be nhwc
    this->outputNames = {"output"};
    this->outputShapes = {oshape};

    readSrcCode(this->clFileName, this->clCode);
    compileKernel(this->clCode, this->clCode.size(), this->clKernelName, this->context, device, this->program, this->kernel);
    prepareMem();
}

cl_int Winograd36To4x4::prepareMem(void) {
    cl_int status;
    std::vector<float> At = {
        1.0000000000f, 1.0000000000f,  1.0000000000f, 1.0000000000f,  1.0000000000f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f, 0.7071067691f, -0.7071067691f, 1.4142135382f, -1.4142135382f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f, 0.4999999702f,  0.4999999702f, 1.9999998808f,  1.9999998808f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f, 0.3535533547f, -0.3535533547f, 2.8284268379f, -2.8284268379f, 1.0000000000f, 0.0000000000f, 0.0000000000f, 
    };
    memObjs.clear();
    memObjs["at"] = createImage2d(context, CL_FLOAT, 8, 1, At.data(), &status);
    
    int inputSize = getTotalSize(inputShapes[0]);
    int outputSize = getTotalSize(outputShapes[0]);
    int biasSize = getTotalSize(inputShapes[1]);
    memObjs["input"] = clCreateBuffer(context, CL_MEM_READ_ONLY |  CL_MEM_ALLOC_HOST_PTR, inputSize*sizeof(float), NULL, &status);
    memObjs["output"] = clCreateBuffer(context, CL_MEM_WRITE_ONLY | CL_MEM_ALLOC_HOST_PTR, outputSize*sizeof(float), NULL, &status);
    memObjs["bias"] = createImage2d(context, CL_FLOAT, biasSize/4, 1, nullptr, &status);

    memTypes["at"] = Image2d;
    memTypes["input"] = Buffer;
    memTypes["output"] = Buffer;
    memTypes["bias"] = Image2d;

    return status;
}

cl_int Winograd36To4x4::setParams(void) {
    // 14400, 4, 64;    32, 4, 2
    globalThreads.resize(3);
    localThreads.resize(3);
    globalThreads[0] = inputShapes[0][2];
    globalThreads[1] = 4;
    globalThreads[2] = inputShapes[0][3] / 4;
    localThreads[0] = 32;
    localThreads[1] = 4;
    localThreads[2] = 2;
    return 0;
}

// cl_int Winograd36To4x4::getOutut(cl_command_queue commandQueue, std::vector<std::vector<float>> &output) {
//     cl_int status;
//     int osize = getTotalSize(this->outputShapes[0]);
//     output.clear();
//     output.resize(1);
//     output[0].resize(osize);
//     status = mslite::clEnqueueReadBuffer(commandQueue, memObjs["output"], CL_TRUE, 0, sizeof(float)*osize, output[0].data(), 0, NULL, NULL);
//     return status;
// }

Winograd36To4x4Naive::Winograd36To4x4Naive(cl_context context, cl_device_id device, std::vector<int> ishape, std::vector<int> oshape): 
    Winograd36To4x4(context, device, std::string("winograd_36_to_4x4_naive"), ishape, oshape) {
}

cl_int Winograd36To4x4Naive::bind(void) {
    cl_int status;

    cl_int4 input_sliced_shape[] = {inputShapes[0][0], inputShapes[0][1], inputShapes[0][2], inputShapes[0][3]/4};
    cl_int4 output_sliced_shape[] = {outputShapes[0][0], outputShapes[0][1], outputShapes[0][2], outputShapes[0][3]/4};

    status = clSetKernelArg(kernel, 0, sizeof(cl_mem), static_cast<void *>(&memObjs["output"]));
    status = clSetKernelArg(kernel, 1, sizeof(cl_mem), static_cast<void *>(&memObjs["input"]));
    status = clSetKernelArg(kernel, 2, sizeof(cl_mem), static_cast<void *>(&memObjs["at"]));
    status = clSetKernelArg(kernel, 3, sizeof(cl_mem), static_cast<void *>(&memObjs["bias"]));
    status = clSetKernelArg(kernel, 4, sizeof(cl_int4), static_cast<void *>(input_sliced_shape));
    status = clSetKernelArg(kernel, 5, sizeof(cl_int4), static_cast<void *>(output_sliced_shape));
    return status;
}

Winograd36To4x4Tflite::Winograd36To4x4Tflite(cl_context context, cl_device_id device, std::vector<int> ishape, std::vector<int> oshape):
    Winograd36To4x4(context, device, std::string("winograd_36_to_4x4_tflite"), ishape, oshape) {

}

cl_int Winograd36To4x4Tflite::bind(void) {
    cl_int status;

    cl_int4 a[] = {outputShapes[0][2]/4, outputShapes[0][2], outputShapes[0][1], outputShapes[0][3]/4};
    cl_int4 b[] = {36, inputShapes[0][2], 0, 0};

    status = clSetKernelArg(kernel, 0, sizeof(cl_mem), static_cast<void *>(&memObjs["output"]));
    status = clSetKernelArg(kernel, 1, sizeof(cl_mem), static_cast<void *>(&memObjs["input"]));
    status = clSetKernelArg(kernel, 2, sizeof(cl_mem), static_cast<void *>(&memObjs["at"]));
    status = clSetKernelArg(kernel, 3, sizeof(cl_mem), static_cast<void *>(&memObjs["bias"]));
    status = clSetKernelArg(kernel, 4, sizeof(cl_int4), static_cast<void *>(a));
    status = clSetKernelArg(kernel, 5, sizeof(cl_int4), static_cast<void *>(b));
    return status;
}

Winograd36To4x4Image2d::Winograd36To4x4Image2d(cl_context context, cl_device_id device, std::vector<int> ishape, std::vector<int> oshape):
    Winograd36To4x4(context, device, std::string("winograd_36_to_4x4_image2d"), ishape, oshape) {

}

cl_int Winograd36To4x4Image2d::prepareMem(void) {
    cl_int status;
    std::vector<float> At = {
        1.0000000000f, 1.0000000000f,  1.0000000000f, 1.0000000000f,  1.0000000000f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f, 0.7071067691f, -0.7071067691f, 1.4142135382f, -1.4142135382f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f, 0.4999999702f,  0.4999999702f, 1.9999998808f,  1.9999998808f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f, 0.3535533547f, -0.3535533547f, 2.8284268379f, -2.8284268379f, 1.0000000000f, 0.0000000000f, 0.0000000000f, 
    };
    memObjs.clear();
    memObjs["at"] = createImage2d(context, CL_FLOAT, 8, 1, At.data(), &status);

    int ih = inputShapes[0][1] * inputShapes[0][3] / 4, iw = inputShapes[0][2];
    int oh = outputShapes[0][1] * outputShapes[0][3] / 4, ow = outputShapes[0][2];
    memObjs["input"] = createImage2d(context, CL_FLOAT, iw, ih, NULL, &status);
    memObjs["output"] = createImage2d(context, CL_FLOAT, ow, oh, NULL, &status);
    
    int biasSize = getTotalSize(inputShapes[1]);
    memObjs["bias"] = clCreateBuffer(context, CL_MEM_READ_ONLY |  CL_MEM_ALLOC_HOST_PTR, biasSize*sizeof(float), NULL, &status);

    memTypes["at"] = Image2d;
    memTypes["input"] = Image2d;
    memTypes["output"] = Image2d;
    memTypes["bias"] = Buffer;

    return status;
}

cl_int Winograd36To4x4Image2d::bind(void) {
    cl_int status;

    cl_int4 input_sliced_shape[] = {inputShapes[0][0], inputShapes[0][1], inputShapes[0][2], inputShapes[0][3]/4};
    cl_int4 output_sliced_shape[] = {outputShapes[0][0], outputShapes[0][1], outputShapes[0][2], outputShapes[0][3]/4};

    status = clSetKernelArg(kernel, 0, sizeof(cl_mem), static_cast<void *>(&memObjs["output"]));
    status = clSetKernelArg(kernel, 1, sizeof(cl_mem), static_cast<void *>(&memObjs["input"]));
    status = clSetKernelArg(kernel, 2, sizeof(cl_mem), static_cast<void *>(&memObjs["at"]));
    status = clSetKernelArg(kernel, 3, sizeof(cl_mem), static_cast<void *>(&memObjs["bias"]));
    status = clSetKernelArg(kernel, 4, sizeof(cl_int4), static_cast<void *>(input_sliced_shape));
    status = clSetKernelArg(kernel, 5, sizeof(cl_int4), static_cast<void *>(output_sliced_shape));
    return status;
}

}